import os
from datetime import datetime
import cv2
import tkinter as tk
from tkinter import messagebox
from PIL import Image, ImageTk
from orbbec.infrared_pipeline import InfraredCamera

pattern = (9, 6)  # 例如 9x6 的棋盘

def detect_chessboard(img, pattern):
    """
    检测棋盘。
    :param img: 相机图像
    :param pattern: 棋盘内角点数量，例如 (9, 6)
    :return: 是否检测到棋盘，检测结果图像
    """
    # 转换为灰度图
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if len(img.shape) == 3 else img

    # 检查图像类型并转换为 8 位灰度图
    if gray.dtype != 'uint8':
        gray = cv2.normalize(gray, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)

    # 检测棋盘角点
    ret, corners = cv2.findChessboardCorners(gray, pattern, flags=cv2.CALIB_CB_ADAPTIVE_THRESH + cv2.CALIB_CB_FAST_CHECK + cv2.CALIB_CB_FILTER_QUADS)

    # 如果检测到棋盘
    if ret:
        # 可视化棋盘角点
        img_show = cv2.drawChessboardCorners(img.copy(), pattern, corners, ret)
        return True, img_show
    else:
        return False, None


def save_or_delete_image(img, img_show, save_path, frame_count, camera_name):
    """
    保存或删除图像。
    :param img: 原始图像
    :param img_show: 检测到棋盘的图像
    :param save_path: 图像保存路径
    :param frame_count: 当前帧计数
    :param camera_name: 相机名称（用于显示和保存）
    :return: 更新后的帧计数
    """
    # 创建 Tkinter 顶层窗口
    top = tk.Toplevel()
    top.title(f"Chessboard - {camera_name}")

    # 将 OpenCV 图像转换为 Tkinter 可显示的格式
    img_show_rgb = cv2.cvtColor(img_show, cv2.COLOR_BGR2RGB)

    # 如果图像是 16 位，将其转换为 8 位
    if img_show_rgb.dtype == 'uint16':
        img_show_rgb = cv2.normalize(img_show_rgb, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U)

    # 将图像转换为 PIL 格式
    img_pil = Image.fromarray(img_show_rgb)
    img_tk = ImageTk.PhotoImage(image=img_pil)

    # 在 Tkinter 窗口中显示图像
    label = tk.Label(top, image=img_tk)
    label.pack()

    # 弹出对话框让用户选择
    response = messagebox.askyesnocancel("检测结果",
                                         f"检测到棋盘！是否保存 {camera_name} 的图像？\n（是：保存，否：删除，取消：继续）")

    if response is True:  # 保存图像
        # 使用 zfill 补零，确保文件名是 8 位
        filename = os.path.join(save_path, f"{str(frame_count).zfill(8)}.png")
        frame_count += 1
        cv2.imwrite(filename, img)
        print(f"Saved frame {frame_count} to {filename}")
    elif response is False:  # 删除图像
        print(f"{camera_name} 图像未保存。")
    else:  # 取消（继续）
        print(f"{camera_name} 图像未保存，继续检测。")

    # 关闭 Tkinter 窗口
    top.destroy()

    return frame_count


def start_recording(mode):
    """
    开始录制。
    :param mode: 录制模式（1: cam1, 2: cam2, 3: both）
    """
    # 创建输出文件夹
    current_time = datetime.now().strftime('%Y%m%d_%H%M%S')
    output_path = f"./output/{current_time}/"
    cam1_path = os.path.join(output_path, "01")
    cam2_path = os.path.join(output_path, "02")
    if not os.path.exists(cam1_path):
        os.makedirs(cam1_path)
    if not os.path.exists(cam2_path):
        os.makedirs(cam2_path)

    # 棋盘内角点数量


    # 初始化红外相机
    infrared_camera = InfraredCamera()
    if infrared_camera.initialize():
        infrared_camera.start()
        try:
            frame_count_cam1 = 0  # cam1 帧计数
            frame_count_cam2 = 0  # cam2 帧计数

            while True:
                # 获取相机数据
                ir_data1 = infrared_camera.get_ir_data(device_index=0)
                ir_data2 = infrared_camera.get_ir_data(device_index=1)
                if ir_data1 is None or ir_data2 is None:
                    continue

                # 转换为 RGB 图像
                ir_image1 = cv2.cvtColor(ir_data1, cv2.COLOR_GRAY2RGB)
                ir_image2 = cv2.cvtColor(ir_data2, cv2.COLOR_GRAY2RGB)

                # 显示图像
                if mode in ['1', '3']:
                    cv2.imshow("Infrared Device1", ir_image1)
                if mode in ['2', '3']:
                    cv2.imshow("Infrared Device2", ir_image2)

                # 检测按键
                key = cv2.waitKey(30)
                # if key != -1:  # 如果有按键按下
                if mode == '1':  # 仅录制 cam1
                    ret1, img_show1 = detect_chessboard(ir_image1, pattern)
                    if ret1:
                        frame_count_cam1 = save_or_delete_image(ir_image1, img_show1, cam1_path, frame_count_cam1,
                                                                "cam1")
                elif mode == '2':  # 仅录制 cam2
                    ret2, img_show2 = detect_chessboard(ir_image2, pattern)
                    if ret2:
                        frame_count_cam2 = save_or_delete_image(ir_image2, img_show2, cam2_path, frame_count_cam2,
                                                                "cam2")
                elif mode == '3':  # 同时录制 cam1 和 cam2
                    ret1, img_show1 = detect_chessboard(ir_image1, pattern)
                    ret2, img_show2 = detect_chessboard(ir_image2, pattern)
                    if ret1 and ret2:  # 两个相机都检测到棋盘
                        frame_count_cam1 = save_or_delete_image(ir_image1, img_show1, cam1_path, frame_count_cam1,
                                                                "cam1")
                        frame_count_cam2 = save_or_delete_image(ir_image2, img_show2, cam2_path, frame_count_cam2,
                                                                "cam2")

                # 按下 'q' 键退出
                if key == ord('q'):
                    break

        except KeyboardInterrupt:
            pass
        finally:
            infrared_camera.stop()
            cv2.destroyAllWindows()


def create_gui():
    """
    创建 GUI 界面。
    """

    def on_mode_selected():
        mode = mode_var.get()
        if mode in ['1', '2', '3']:
            start_recording(mode)
        else:
            messagebox.showerror("错误", "请选择有效的录制模式！")

    # 创建主窗口
    root = tk.Tk()
    root.title("棋盘检测录制程序")

    # 添加模式选择标签
    tk.Label(root, text="请选择录制模式：").pack(pady=10)

    # 添加模式选择单选按钮
    mode_var = tk.StringVar(value="1")
    tk.Radiobutton(root, text="仅录制 cam1", variable=mode_var, value="1").pack(anchor=tk.W)
    tk.Radiobutton(root, text="仅录制 cam2", variable=mode_var, value="2").pack(anchor=tk.W)
    tk.Radiobutton(root, text="同时录制 cam1 和 cam2", variable=mode_var, value="3").pack(anchor=tk.W)

    # 添加开始按钮
    tk.Button(root, text="开始录制", command=on_mode_selected).pack(pady=20)

    # 运行主循环
    root.mainloop()


if __name__ == "__main__":
    create_gui()